#!/bin/bash

# 资产管理系统 - 数据恢复脚本
# 版本: 1.0

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

# 配置变量
APP_DIR="/opt/asset-management"
BACKUP_DIR="/opt/asset_backups"

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 显示使用方法
show_usage() {
    echo "使用方法: $0 <backup_file>"
    echo "示例: $0 /opt/asset_backups/backup_20250807_120000.tar.gz"
    echo
    echo "可用的备份文件:"
    if [[ -d "$BACKUP_DIR" ]]; then
        ls -la "$BACKUP_DIR"/backup_*.tar.gz 2>/dev/null || echo "  没有找到备份文件"
    else
        echo "  备份目录不存在"
    fi
}

# 检查参数
check_parameters() {
    if [[ $# -ne 1 ]]; then
        show_usage
        exit 1
    fi
    
    local backup_file="$1"
    
    if [[ ! -f "$backup_file" ]]; then
        log_error "备份文件不存在: $backup_file"
        exit 1
    fi
    
    BACKUP_FILE="$backup_file"
}

# 检查环境
check_environment() {
    if [[ ! -d "$APP_DIR" ]]; then
        log_error "应用目录不存在: $APP_DIR"
        exit 1
    fi
    
    cd "$APP_DIR"
    
    # 检查Docker服务
    if ! docker-compose ps >/dev/null 2>&1; then
        log_error "Docker服务未运行或配置错误"
        exit 1
    fi
}

# 确认恢复操作
confirm_restore() {
    log_warning "警告: 此操作将覆盖当前的所有数据！"
    log_warning "备份文件: $BACKUP_FILE"
    log_warning "恢复时间: $(date)"
    echo
    read -p "确定要继续恢复操作吗？(输入 'YES' 确认): " confirmation
    
    if [[ "$confirmation" != "YES" ]]; then
        log_info "恢复操作已取消"
        exit 0
    fi
}

# 停止服务
stop_services() {
    log_info "停止应用服务..."
    
    # 停止Web服务但保持数据库运行
    docker-compose stop web nginx celery celery-beat
    
    log_success "应用服务已停止"
}

# 解压备份文件
extract_backup() {
    log_info "解压备份文件..."
    
    local temp_dir="/tmp/asset_restore_$(date +%s)"
    mkdir -p "$temp_dir"
    
    tar -xzf "$BACKUP_FILE" -C "$temp_dir"
    
    # 查找解压后的目录
    RESTORE_DIR=$(find "$temp_dir" -maxdepth 1 -type d -name "20*" | head -1)
    
    if [[ -z "$RESTORE_DIR" ]]; then
        log_error "备份文件格式错误"
        rm -rf "$temp_dir"
        exit 1
    fi
    
    log_success "备份文件解压完成: $RESTORE_DIR"
}

# 恢复数据库
restore_database() {
    log_info "恢复数据库..."
    
    local db_file="$RESTORE_DIR/database.sql.gz"
    
    if [[ ! -f "$db_file" ]]; then
        log_error "数据库备份文件不存在: $db_file"
        exit 1
    fi
    
    # 删除现有数据库
    log_info "删除现有数据库..."
    docker-compose exec -T db mysql -u root -p123456 -e "DROP DATABASE IF EXISTS asset_db;"
    docker-compose exec -T db mysql -u root -p123456 -e "CREATE DATABASE asset_db CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;"
    
    # 恢复数据库
    log_info "导入数据库备份..."
    zcat "$db_file" | docker-compose exec -T db mysql -u root -p123456 asset_db
    
    log_success "数据库恢复完成"
}

# 恢复媒体文件
restore_media() {
    log_info "恢复媒体文件..."
    
    # 备份当前媒体文件
    local current_media_backup="/tmp/current_media_$(date +%s)"
    docker cp asset_web:/app/media "$current_media_backup" 2>/dev/null || true
    
    # 删除当前媒体文件
    docker-compose exec -T web rm -rf /app/media/* 2>/dev/null || true
    
    # 恢复媒体文件
    if [[ -d "$RESTORE_DIR/media" ]]; then
        docker cp "$RESTORE_DIR/media/." asset_web:/app/media/
        log_success "媒体文件恢复完成"
    else
        log_warning "备份中没有媒体文件"
    fi
    
    # 恢复静态文件
    if [[ -d "$RESTORE_DIR/static" ]]; then
        docker cp "$RESTORE_DIR/static/." asset_web:/app/static/
        log_success "静态文件恢复完成"
    else
        log_warning "备份中没有静态文件"
    fi
}

# 恢复配置文件
restore_config() {
    log_info "恢复配置文件..."
    
    # 备份当前配置
    cp .env .env.backup.$(date +%s) 2>/dev/null || true
    
    # 恢复环境配置
    if [[ -f "$RESTORE_DIR/.env" ]]; then
        cp "$RESTORE_DIR/.env" .env
        log_success "环境配置恢复完成"
    else
        log_warning "备份中没有环境配置文件"
    fi
    
    # 恢复Docker配置
    if [[ -d "$RESTORE_DIR/docker" ]]; then
        cp -r "$RESTORE_DIR/docker" ./
        log_success "Docker配置恢复完成"
    else
        log_warning "备份中没有Docker配置"
    fi
}

# 启动服务
start_services() {
    log_info "启动应用服务..."
    
    # 重新构建镜像（如果配置有变化）
    docker-compose build --no-cache web
    
    # 启动所有服务
    docker-compose up -d
    
    # 等待服务启动
    sleep 30
    
    log_success "应用服务已启动"
}

# 验证恢复
verify_restore() {
    log_info "验证恢复结果..."
    
    # 检查数据库连接
    if docker-compose exec -T web python manage.py check --database default >/dev/null 2>&1; then
        log_success "数据库连接正常"
    else
        log_error "数据库连接失败"
        return 1
    fi
    
    # 检查Web服务
    local max_attempts=10
    local attempt=1
    
    while [[ $attempt -le $max_attempts ]]; do
        if curl -s -o /dev/null -w "%{http_code}" http://localhost/health/ | grep -q "200"; then
            log_success "Web服务运行正常"
            break
        fi
        log_info "等待Web服务启动... ($attempt/$max_attempts)"
        sleep 10
        ((attempt++))
    done
    
    if [[ $attempt -gt $max_attempts ]]; then
        log_warning "Web服务健康检查超时"
        return 1
    fi
    
    return 0
}

# 清理临时文件
cleanup() {
    log_info "清理临时文件..."
    
    if [[ -n "$RESTORE_DIR" && -d "$(dirname "$RESTORE_DIR")" ]]; then
        rm -rf "$(dirname "$RESTORE_DIR")"
    fi
    
    log_success "临时文件清理完成"
}

# 显示恢复信息
show_restore_info() {
    log_success "=== 数据恢复完成 ==="
    echo
    log_info "恢复信息:"
    log_info "  备份文件: $BACKUP_FILE"
    log_info "  恢复时间: $(date)"
    echo
    log_info "服务状态:"
    docker-compose ps
    echo
    log_warning "重要提醒:"
    log_warning "1. 请检查应用功能是否正常"
    log_warning "2. 请验证数据完整性"
    log_warning "3. 建议立即创建新的备份"
}

# 错误处理
error_handler() {
    log_error "恢复过程中发生错误，正在清理..."
    cleanup
    
    # 尝试重启服务
    log_info "尝试重启服务..."
    docker-compose restart 2>/dev/null || true
    
    exit 1
}

# 主函数
main() {
    log_info "开始执行数据恢复..."
    log_info "恢复时间: $(date)"
    
    # 设置错误处理
    trap error_handler ERR
    
    check_parameters "$@"
    check_environment
    confirm_restore
    stop_services
    extract_backup
    restore_database
    restore_media
    restore_config
    start_services
    
    if verify_restore; then
        cleanup
        show_restore_info
        log_success "数据恢复完成！"
    else
        log_error "恢复验证失败，请检查系统状态"
        exit 1
    fi
}

# 执行主函数
main "$@"